import os

import cv2
import numpy as np


class DigitRecognizer:
    def __init__(self, template_folder="templates"):
        """
        初始化数字识别器，加载模板图片
        :param template_folder: 模板文件夹路径，假设模板文件命名为 '0.png', '1.png', ..., '9.png'
        """
        self.template_folder = template_folder
        self.templates = self.load_templates()

    def load_templates(self):
        """
        从指定文件夹加载模板图片，假设模板文件命名为 '0.png', '1.png', ..., '9.png'
        返回一个字典，键为数字，值为对应的灰度图模板
        """
        templates = {}
        self.script_dir = os.path.dirname(os.path.abspath(__file__))
        for digit in range(10):
            path = os.path.join(self.script_dir, f"{self.template_folder}", f"{digit}.png")
            temp = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
            if temp is None:
                print(f"模板 {path} 加载失败")
                continue
            # 可根据需要对模板图像进行预处理，例如二值化、调整大小等
            _, temp = cv2.threshold(temp, 127, 255, cv2.THRESH_BINARY)
            templates[digit] = temp
        return templates

    def recognize_digit(self, input_image):
        """
        利用模板匹配识别输入的单个数字图像
        :param input_image: 待识别的数字图像（灰度图）
        :return: 最佳匹配的数字和对应匹配得分
        """
        best_digit = None
        best_score = -np.inf  # 对于TM_CCOEFF_NORMED，得分越高越好

        # 对输入图像进行预处理（如二值化）
        _, input_bin = cv2.threshold(input_image, 127, 255, cv2.THRESH_BINARY)

        for digit, template in self.templates.items():
            # 如果模板大小与输入图像不同，可以先进行resize或者局部匹配
            res = cv2.matchTemplate(input_bin, template, cv2.TM_CCOEFF_NORMED)
            # 获取最大匹配值
            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
            # 打印每个模板的匹配得分，便于调试
            # print(f"Digit: {digit} score = {max_val:.3f}")
            if max_val > best_score:
                best_score = max_val
                best_digit = digit
        # print(f"=== Digit: {best_digit} score = {best_score:.3f} ===")

        return best_digit, best_score


if __name__ == '__main__':
    """
    读取模板中的图像，标记为0-9
    拿到传入的已经切割好的小图片依次对比图像，找到最匹配的数字
    """
    # 创建DigitRecognizer实例并传入模板文件夹路径
    debug = DigitRecognizer()

    # 参数为灰度图
    debug.recognize_digit()
